/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution
{
public:

    // 直接查询，超时
    vector<int> treeQueries(TreeNode *root, vector<int> &queries)
    {
        unordered_map<TreeNode *, int> subTreeHeight;
        unordered_map<int, TreeNode *> valNode;
        unordered_map<TreeNode *, TreeNode *> parentMap;
        setupMaps(root, nullptr, subTreeHeight, valNode, parentMap);
        vector<int> result;
        for (int query : queries)
        {
            TreeNode *t = valNode[query];
            int subHeight = subTreeHeight[t];
            t = parentMap[t];
            {
                int leftSubHeight = 0;
                int rightSubHeight = 0;
                if (t->left != nullptr)
                {
                    leftSubHeight = subTreeHeight[t->left];
                }
                if (t->right != nullptr)
                {
                    rightSubHeight = subTreeHeight[t->right];
                }
                // subHeight一定是leftSubHeight，rightSubHeight之一，删掉后，父节点的高度就变为另一个+1
                if (leftSubHeight == subHeight)
                {
                    subHeight = rightSubHeight + 1;
                }
                else
                {
                    subHeight = leftSubHeight + 1;
                }

                if (subTreeHeight[t] == subHeight)
                {
                    result.push_back(subTreeHeight[root] - 1);
                    continue;
                }
                else
                {
                    subTreeHeight[t] = subHeight;
                }
                t = parentMap[t];
            }
            while (t != nullptr) // 向上更新子树高度
            {
                int leftSubHeight = 0;
                int rightSubHeight = 0;
                if (t->left != nullptr)
                {
                    leftSubHeight = subTreeHeight[t->left];
                }
                if (t->right != nullptr)
                {
                    rightSubHeight = subTreeHeight[t->right];
                }
                int newSubHeight = max(leftSubHeight, rightSubHeight) + 1;
                if (subTreeHeight[t] == newSubHeight)
                {
                    break;
                }
                else
                {
                    subTreeHeight[t] = newSubHeight;
                }
                t = parentMap[t];
            }
            result.push_back(subTreeHeight[root] - 1);

            // 变回去
            t = valNode[query];
            t = parentMap[t];
            while (t != nullptr) // 向上更新子树高度
            {
                int leftSubHeight = 0;
                int rightSubHeight = 0;
                if (t->left != nullptr)
                {
                    leftSubHeight = subTreeHeight[t->left];
                }
                if (t->right != nullptr)
                {
                    rightSubHeight = subTreeHeight[t->right];
                }
                int newSubHeight = max(leftSubHeight, rightSubHeight) + 1;
                if (subTreeHeight[t] == newSubHeight)
                {
                    break;
                }
                else
                {
                    subTreeHeight[t] = newSubHeight;
                }
                t = parentMap[t];
            }
        }
        return result;
    }

    void setupMaps(TreeNode *t, TreeNode *parent,
                   unordered_map<TreeNode *, int> &heightMap,
                   unordered_map<int, TreeNode *> &valMap,
                   unordered_map<TreeNode *, TreeNode *> &parentMap)
    {
        parentMap[t] = parent;
        valMap[t->val] = t;
        if (t->left == nullptr && t->right == nullptr)
        {
            heightMap[t] = 1;
            return;
        }
        int height = 1;
        if (t->left != nullptr)
        {
            setupMaps(t->left, t, heightMap, valMap, parentMap);
            height = max(height, heightMap[t->left] + 1);
        }
        if (t->right != nullptr)
        {
            setupMaps(t->right, t, heightMap, valMap, parentMap);
            height = max(height, heightMap[t->right] + 1);
        }
        heightMap[t] = height;
    }
};